# 导入必要的模块
from torch import nn  # PyTorch 的神经网络模块，用于定义网络层
from transformers import AutoConfig  # 用于加载预训练模型的配置
from transformers import AutoTokenizer  # 用于加载预训练模型的分词器
import torch  # PyTorch，用于张量操作
import torch.nn.functional as F  # PyTorch 的函数式 API，包含 softmax 等操作
from math import sqrt  # 导入数学库中的平方根函数，用于缩放操作




# 定义 Scaled Dot-Product Attention（缩放点积注意力）函数
def scaled_dot_product_attention(query, key, value, query_mask=None, key_mask=None, mask=None):
    """
    实现缩放点积注意力机制（Scaled Dot-Product Attention）。

    参数：
    - query: 查询张量，形状为 [batch_size, seq_len_q, dim_k]
    - key: 键张量，形状为 [batch_size, seq_len_k, dim_k]
    - value: 值张量，形状为 [batch_size, seq_len_k, dim_v]
    - query_mask: 查询的可选掩码张量，形状为 [batch_size, seq_len_q]，用于标记有效的查询位置
    - key_mask: 键的可选掩码张量，形状为 [batch_size, seq_len_k]，用于标记有效的键位置
    - mask: 自定义的掩码张量，形状为 [batch_size, seq_len_q, seq_len_k]，用于屏蔽特定的注意力分数

    返回：
    - 注意力输出张量，形状为 [batch_size, seq_len_q, dim_v]
    """
    # 获取键张量的最后一维大小（dim_k），即隐藏维度
    dim_k = query.size(-1)

    # 计算点积注意力分数（未缩放），形状为 [batch_size, seq_len_q, seq_len_k]
    # torch.bmm 是批量矩阵乘法，query 和 key 的转置相乘
    scores = torch.bmm(query, key.transpose(1, 2)) / sqrt(dim_k)  # 缩放分数，防止数值过大导致梯度消失

    # 如果提供了查询掩码和键掩码，则生成联合掩码
    if query_mask is not None and key_mask is not None:
        # query_mask 的形状扩展为 [batch_size, seq_len_q, 1]
        # key_mask 的形状扩展为 [batch_size, 1, seq_len_k]
        # 两者相乘后，生成形状为 [batch_size, seq_len_q, seq_len_k] 的联合掩码
        mask = torch.bmm(query_mask.unsqueeze(-1), key_mask.unsqueeze(1))

        # 如果提供了掩码（mask），则将被屏蔽的位置的分数设置为负无穷（-inf）
    # 这样在 softmax 中，这些位置的权重会变为 0
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -float("inf"))

        # 对注意力分数进行 softmax 操作，得到注意力权重
    # softmax 在最后一维（seq_len_k）上计算，形状为 [batch_size, seq_len_q, seq_len_k]
    weights = F.softmax(scores, dim=-1)

    # 使用注意力权重加权值张量，得到最终的注意力输出
    # torch.bmm(weights, value) 的形状为 [batch_size, seq_len_q, dim_v]
    return torch.bmm(weights, value) #torch.bmm 是批量矩阵乘法，query 和 key 的转置相乘

# 定义一个注意力头（Attention Head）类，继承自 PyTorch 的 nn.Module
class AttentionHead(nn.Module):
    def __init__(self, embed_dim, head_dim):
        """
        初始化注意力头。

        参数：
        - embed_dim: 输入嵌入的维度（embedding dimension），即输入的特征维度。
        - head_dim: 每个注意力头的维度（head dimension），即查询、键和值的输出特征维度。
        """
        super().__init__()  # 调用父类的初始化方法

        # 定义线性变换，用于生成 Query、Key 和 Value
        # nn.Linear(embed_dim, head_dim) 表示将输入从 embed_dim 映射到 head_dim
        self.q = nn.Linear(embed_dim, head_dim)  # 用于生成 Query 的线性层
        self.k = nn.Linear(embed_dim, head_dim)  # 用于生成 Key 的线性层
        self.v = nn.Linear(embed_dim, head_dim)  # 用于生成 Value 的线性层

    def forward(self, query, key, value, query_mask=None, key_mask=None, mask=None):
        """
        前向传播函数，计算注意力输出。

        参数：
        - query: 查询张量，形状为 [batch_size, seq_len_q, embed_dim]
        - key: 键张量，形状为 [batch_size, seq_len_k, embed_dim]
        - value: 值张量，形状为 [batch_size, seq_len_k, embed_dim]
        - query_mask: 查询的可选掩码张量，形状为 [batch_size, seq_len_q]，用于标记有效的查询位置。
        - key_mask: 键的可选掩码张量，形状为 [batch_size, seq_len_k]，用于标记有效的键位置。
        - mask: 自定义的掩码张量，形状为 [batch_size, seq_len_q, seq_len_k]，用于屏蔽特定的注意力分数。

        返回：
        - attn_outputs: 注意力输出张量，形状为 [batch_size, seq_len_q, head_dim]
        """
        # 通过线性层生成 Query、Key 和 Value
        # self.q(query): 将 query 输入映射到 head_dim 维度，形状为 [batch_size, seq_len_q, head_dim]
        # self.k(key): 将 key 输入映射到 head_dim 维度，形状为 [batch_size, seq_len_k, head_dim]
        # self.v(value): 将 value 输入映射到 head_dim 维度，形状为 [batch_size, seq_len_k, head_dim]
        attn_outputs = scaled_dot_product_attention(
            self.q(query),  # 生成 Query
            self.k(key),    # 生成 Key
            self.v(value),  # 生成 Value
            query_mask,     # 查询掩码
            key_mask,       # 键掩码
            mask            # 自定义掩码
        )

        # 返回注意力输出，形状为 [batch_size, seq_len_q, head_dim]
        return attn_outputs
# 定义多头注意力（Multi-Head Attention）类，继承自 PyTorch 的 nn.Module
class MultiHeadAttention(nn.Module):
    def __init__(self, config):
        """
        初始化多头注意力模块。

        参数：
        - config: 配置对象，包含以下属性：
            - hidden_size: 输入嵌入的维度（embedding dimension），即输入的特征维度。
            - num_attention_heads: 注意力头的数量（number of attention heads）。
        """
        super().__init__()  # 调用父类的初始化方法

        # 从配置中获取嵌入维度和注意力头的数量
        embed_dim = config.hidden_size  # 输入嵌入维度
        num_heads = config.num_attention_heads  # 注意力头的数量

        # 每个注意力头的维度（head_dim），通过将总嵌入维度除以注意力头数量计算
        head_dim = embed_dim // num_heads

        # 定义多个注意力头，每个注意力头是一个 AttentionHead 实例
        # 使用 nn.ModuleList 存储多个注意力头
        self.heads = nn.ModuleList(
            [AttentionHead(embed_dim, head_dim) for _ in range(num_heads)]
        )

        # 定义一个线性层，用于将多头注意力的输出映射回原始嵌入维度
        self.output_linear = nn.Linear(embed_dim, embed_dim)

    def forward(self, query, key, value, query_mask=None, key_mask=None, mask=None):
        """
        前向传播函数，计算多头注意力的输出。

        参数：
        - query: 查询张量，形状为 [batch_size, seq_len_q, embed_dim]
        - key: 键张量，形状为 [batch_size, seq_len_k, embed_dim]
        - value: 值张量，形状为 [batch_size, seq_len_k, embed_dim]
        - query_mask: 查询的可选掩码张量，形状为 [batch_size, seq_len_q]，用于标记有效的查询位置。
        - key_mask: 键的可选掩码张量，形状为 [batch_size, seq_len_k]，用于标记有效的键位置。
        - mask: 自定义的掩码张量，形状为 [batch_size, seq_len_q, seq_len_k]，用于屏蔽特定的注意力分数。

        返回：
        - x: 多头注意力的输出张量，形状为 [batch_size, seq_len_q, embed_dim]
        """
        # 计算每个注意力头的输出，并将它们在最后一维（head_dim）上拼接
        # 对每个 AttentionHead 实例调用 forward 方法
        # 每个头的输出形状为 [batch_size, seq_len_q, head_dim]
        # 拼接后的形状为 [batch_size, seq_len_q, embed_dim]
        x = torch.cat([
            h(query, key, value, query_mask, key_mask, mask) for h in self.heads
        ], dim=-1)

        # 通过线性层将拼接后的多头注意力输出映射回原始嵌入维度
        # 输出形状仍为 [batch_size, seq_len_q, embed_dim]
        x = self.output_linear(x)

        # 返回多头注意力的最终输出
        return x

if __name__ == '__main__':

    model_ckpt = "bert-base-uncased"

    # 加载与指定模型对应的分词器
    tokenizer = AutoTokenizer.from_pretrained(model_ckpt)

    # 定义输入文本
    text = "time flies like an arrow"

    # 使用分词器将文本转为 token ID
    # 参数说明：
    # - return_tensors="pt"：返回 PyTorch 张量格式
    # - add_special_tokens=False：不添加特殊标记（如 [CLS] 和 [SEP]）,特殊标记的作用总结
    # [CLS]：句子级别的表示，用于分类任务。
    # [SEP]：分隔句子或文本片段，用于句子对任务。
    # [PAD]：对齐序列长度，用于批量处理。
    # [MASK]：掩盖 token，用于预训练任务。
    inputs = tokenizer(text, return_tensors="pt", add_special_tokens=False)

    # 打印分词后的 token ID
    # 例如，输出可能是：tensor([[2051, 21695, 2066, 2019, 8617]])
    print(inputs.input_ids) #tensor([[ 2051, 10029,  2066,  2019,  8612]])

    # 加载预训练模型的配置
    # AutoConfig 会从模型检查点中加载配置文件，包含模型的超参数信息
    config = AutoConfig.from_pretrained(model_ckpt)

    # 定义一个嵌入层（Embedding Layer）
    # nn.Embedding 是一个 嵌入层（Embedding Layer），它的作用是将输入的 token ID 转换为对应的 词向量（word embeddings）
    # 参数：
    # - config.vocab_size：词汇表大小（BERT 的词汇表通常是 30522）
    # - config.hidden_size：隐藏层维度（BERT 的默认隐藏层维度是 768）
    token_emb = nn.Embedding(config.vocab_size, config.hidden_size) #初始化embedding class

    # 打印嵌入层的结构信息
    # 例如：Embedding(30522, 768)
    print(token_emb)

    # 将分词后的 token ID 输入到嵌入层，生成对应的词向量
    # inputs.input_ids 的形状为 [batch_size, sequence_length]
    # 生成的 inputs_embeds 的形状为 [batch_size, sequence_length, hidden_size]
    inputs_embeds = token_emb(inputs.input_ids) #使用embedding

    # 打印词向量的张量形状
    # 例如：torch.Size([1, 5, 768])，表示 1 个样本，5 个 token，每个 token 的向量维度为 768
    print(inputs_embeds.shape)

    multihead_attn = MultiHeadAttention(config)
    query = key = value = inputs_embeds
    attn_output = multihead_attn(query, key, value)
    print(attn_output)
    print(attn_output.shape)